package com.demo.component.ollama.generate.impl;

import com.demo.component.ollama.generate.inter.GenerateSync;
import io.github.ollama4j.OllamaAPI;
import io.github.ollama4j.exceptions.OllamaBaseException;
import io.github.ollama4j.models.embeddings.OllamaEmbedRequestModel;
import io.github.ollama4j.models.embeddings.OllamaEmbedResponseModel;
import io.github.ollama4j.models.response.OllamaResult;
import io.github.ollama4j.types.OllamaModelType;
import io.github.ollama4j.utils.OptionsBuilder;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

import java.io.File;
import java.io.IOException;
import java.util.Arrays;
import java.util.List;

@Service("GenerateSyncImpl")
public class GenerateSyncImpl implements GenerateSync {

    @Value("${application.ollama.url}")
    private String url;

    @Value("${application.ollama.model}")
    private String model;

    @Override
    public String generateAnswer(String prompt) {

        OllamaAPI ollamaAPI = new OllamaAPI(url);

        try {
            OllamaResult result =
                    ollamaAPI.generate(model, prompt, null);

            return result.getResponse();
        } catch (OllamaBaseException e) {
            throw new RuntimeException(e);
        } catch (IOException e) {
            throw new RuntimeException(e);
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public String generateFromImage(String filePath, String prompt) {

        OllamaAPI ollamaAPI = new OllamaAPI(url);

        try {
            File imageFile = new File(filePath);

            OllamaResult result = ollamaAPI.generateWithImageFiles(OllamaModelType.LLAVA,
                    prompt, List.of(imageFile),
                    new OptionsBuilder().build()
            );

            return result.getResponse();
        } catch (OllamaBaseException e) {
            throw new RuntimeException(e);
        } catch (IOException e) {
            throw new RuntimeException(e);
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public List<List<Double>> generateEmbeddings(String prompt) {

        OllamaAPI ollamaAPI = new OllamaAPI(url);

        try {
            OllamaEmbedRequestModel requestModel = new OllamaEmbedRequestModel(model, Arrays.asList(prompt));

            OllamaEmbedResponseModel responseModel = ollamaAPI.embed(requestModel);

            return responseModel.getEmbeddings();
        } catch (OllamaBaseException e) {
            throw new RuntimeException(e);
        } catch (IOException e) {
            throw new RuntimeException(e);
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
    }
}
